#include "CoderImp.h"
#include <algorithm>
#include <iostream>
using namespace std;

bool operator== (const DictData& lf, const DictData& rf)
{
	if(lf._dataSize == rf._dataSize)
	{
		int i = 0;
		while(i < lf._dataSize && lf._data[i] == rf._data[i])
			++i;
		return i == lf._dataSize;
	}
	return false;
}

bool operator< (const DictData& lf, const DictData& rf)
{
	if(lf._dataSize == rf._dataSize)
	{
		int i = 0;
		while(i < lf._dataSize && lf._data[i] == rf._data[i])
			++i;
		return (i != lf._dataSize && lf._data[i] < rf._data[i]);
	}
	else if(lf._dataSize < rf._dataSize)
		return true;

	return false;
}

bool operator< (const DictItemExt& lf, const DictItemExt& rf)
{
	return lf._count < rf._count;
}

#include <stack>

class HuffmanTreeVisitor : protected stack<HuffmanNode*>
{
public:
	HuffmanTreeVisitor() { _code = 0; _codeSize = 0; }

// left node: isLeft==0, right node: isLeft==1, root node: isLeft==-1
	void Push(HuffmanNode* node, int isLeft)
	{
		this->push(node);
		_code = _code << 1;
		++_codeSize;
		if(isLeft != -1)
			_code += isLeft;
	}

	HuffmanNode* Pop()
	{
		HuffmanNode* ret = this->top();
		if(ret && ret->_itemExt)
		{
			ret->_itemExt->_item._code._code = _code;
			ret->_itemExt->_item._code._bitSize = _codeSize;
		//	printf("%x     ", ret->_itemExt->_item._data._data);
		//	printf("%x, %d\n", _code, this->size()-1);
		}
		this->pop();
		_code = _code >> 1;
		--_codeSize;
		return ret;
	}

	bool IsStackEmpty()	{ return this->empty(); }
	HuffmanNode* Top() { return this->top(); }

private:
	int _code;
	int _codeSize;
};

void VisitTree(HuffmanNode* t, int code, int deep)
{
	if(t)
	{
		if(t->_itemExt)
		{
			t->_itemExt->_item._code._code = code;
			t->_itemExt->_item._code._bitSize = deep;
		}
		VisitTree(t->left, (code<<1)+0, deep+1);
		VisitTree(t->right, (code<<1)+1, deep+1);

		delete t;
	}
}

//-----------------------------------------------------------------------------

CoderImp::CoderImp(void) : 
	_head(0), 
	_cur(0),
	_dictItemCount(0),
	_dictItemAllCount(0)
{
}


CoderImp::~CoderImp(void)
{
	for(vector<DictItemExt>::iterator it = _dictionary.begin(); it != _dictionary.end(); ++it)
	{
		delete it->_item._data._data;
	}
}

void CoderImp::Input(BYTE data)
{
	DictData dd;
	dd._dataSize = 1;
	dd._data = new BYTE[1];
	dd._data[0] = data;
	++_counter[dd];
	++_dictItemAllCount;
}

void CoderImp::Parse()
{
	map<DictData, int>::iterator mapIt;
	for (mapIt = _counter.begin(); mapIt != _counter.end(); ++mapIt)
	{
		DictItemExt dn;
		dn._item._data._data = mapIt->first._data;
		dn._item._data._dataSize = mapIt->first._dataSize;
		dn._count = mapIt->second;
		_dictionary.push_back(dn);
	}

	sort(_dictionary.begin(), _dictionary.end());

// To make the list, we see every list node as a 'root', 
// so rootNum will increase 1 every time.
	int rootNum = 0;
	for(vector<DictItemExt>::iterator it = _dictionary.begin(); 
			it != _dictionary.end(); ++it, ++rootNum)
	{
		if(0 == _head)
		{
			_head = new HuffmanNode();
			_cur = _head;
		}
		else
		{
			_cur->next = new HuffmanNode();
			_cur = _cur->next;
		}
		_cur->_itemExt = &(*it);
		_cur->count = _cur->_itemExt->_count;
	}

// To make the list become a Binary Tree. 
// The root number will decrease til it is two. 
// In fact the original list become two binary tree,
// we'll do the last thing to merge them to one tree.
	while(rootNum > 2)
	{
		HuffmanNode* newNode = new HuffmanNode();
		newNode->left = _head;
		newNode->right = _head->next;
		newNode->count = newNode->left->count + newNode->right->count;
		newNode->next = _head->next->next;
		_head = newNode;
		--rootNum; // the root number is decreased by up-operation.

		while(newNode->next !=0 && newNode->count > newNode->next->count)
		{
			HuffmanNode* tmp;
			if(newNode == _head)
			{
				tmp = newNode->next;
				newNode->next = tmp->next;
				_head = tmp;
				tmp->next = newNode;
				_cur = _head;
			}
			else
			{
				tmp = newNode->next;
				newNode->next = tmp->next;
				_cur->next = tmp;
				tmp->next = newNode;
				_cur = tmp;
			}
		}
	}
	HuffmanNode* root = new HuffmanNode();
	root->left = _head;
	root->right = _head->next;

// Now, the binary tree is constructed. 
// We can compute the huffman code by visit the tree.
// ! Note that we'd delete the pointer after visited tree-nodes to avoid leaks.
	//HuffmanTreeVisitor visitor;
	//visitor.Push(root, -1);
	//while(!visitor.IsStackEmpty())
	//{
	//	HuffmanNode* top;
	//	while(top = visitor.Top(), top)
	//	{
	//		visitor.Push(top->left, 0);
	//	}
	//	visitor.Pop();
	//	if(!visitor.IsStackEmpty())
	//	{
	//		visitor.Push(visitor.Pop()->right, 1);
	//	}
	//}
	VisitTree(root, 0, 0); // Recursion method, it is not good.

// Construct a easy-qeury map.
	for(vector<DictItemExt>::iterator it = _dictionary.begin();
		it != _dictionary.end(); ++it)
	{
		_outmap[it->_item._data] = it->_item._code;
	}
}

int CoderImp::GetDataCount()
{
	return _dictItemAllCount;
}

int CoderImp::GetDictItemCount()
{
	return _dictionary.size();
}

void CoderImp::InputBlock(BYTE* data, int dataSize)
{
	_dataBlock = data;
	_dataBlockSize = dataSize;
	_dataIter = 0;
}

bool CoderImp::GetNextDictItem(DictItem* di)
{
	if(_dictItemCount == _dictionary.size())
		return false;
	di->_code = _dictionary[_dictItemCount]._item._code;
	di->_data = _dictionary[_dictItemCount]._item._data;
	_dictItemCount += 1;
	return true;
}

bool CoderImp::GetNextCode(DataCode* oc)
{
	if(_dataIter >= _dataBlockSize)
		return false;

	DictData dd;
	dd._data = _dataBlock + _dataIter;
	dd._dataSize = 1;
	_dataIter += 1;
	
	oc->_bitSize = _outmap[dd]._bitSize;
	oc->_code = _outmap[dd]._code;
	
	return true;
}
